import torch
import torch.nn as nn
from torchvision.models import resnet152

ckptpath = '/mnt/md1/Downloads/model/resnet152_80e_71.38p.pth.tar'
D = torch.load(ckptpath)
model = resnet152()
D.keys()
state_dict = D['state_dict']

# 根据需要改load函数
def try_load_state_dict(self, state_dict, strict=True):
    from torch.nn.parameter import Parameter
    own_state = self.state_dict()
    for name, param in state_dict.items():
        newname = name.replace('module.base_model.','')
        if newname in own_state:
            if isinstance(param, Parameter):
                # backwards compatibility for serialized parameters
                param = param.data
            try:
                own_state[newname].copy_(param)
            except Exception:
                raise RuntimeError('While copying the parameter named {}, '
                                   'whose dimensions in the model are {} and '
                                   'whose dimensions in the checkpoint are {}.'
                                   .format(name, own_state[name].size(), param.size()))
        elif strict:
            raise KeyError('unexpected key "{}" in state_dict'
                           .format(name))
    if strict:
        missing = set(own_state.keys()) - set(state_dict.keys())
        if len(missing) > 0:
            raise KeyError('missing keys in state_dict: "{}"'.format(missing))

# 提特征
model.try_load_state_dict = try_load_state_dict
model.try_load_state_dict(model,state_dict,strict=False)

# 调整最后一层
model.fc = nn.Linear(2048,201)
# 删除多余的函数
del model.try_load_state_dict

# 检查是否正确
# model.state_dict().update(model_weight)
# model.load_state_dict()
# id(model.state_dict())
# id(model_weight)
# model.state_dict()['layer4.2.bn3.bias']
# model_weight['layer4.2.bn3.bias']
# state_dict['module.base_model.layer4.2.bn3.bias']

# save model
torch.save(model,'/mnt/md1/Model/raw_resnet152_kinitics400_rgb_model.pth')
torch.save(model.state_dict(),'/mnt/md1/Model/raw_resnet152_kinitics400_rgb_staticdict.pth')

# 模型的输入参数和训练和使用示例
import torchvision.transforms as transforms
from torch.autograd import Variable
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

model = torch.load('/mnt/md1/Model/raw_resnet152_kinitics400_rgb_model.pth')

I = torch.randn(8,3,224,224)
I = Variable(I)
O = model(I)
O.shape

for item in model.childrens():
    print(item)
